
import cv2
import numpy as np
import paho.mqtt.client as mqtt
import base64
import time
import random
import json
import threading
import platform
import os
from datetime import datetime
import pytz
import mss
from PIL import Image
import getpass
import hashlib

# 导入配置
from config import MQTT_CONFIG, SCREENSHOT_CONFIG

# MQTT 配置
broker = MQTT_CONFIG['broker']
port = MQTT_CONFIG['port']
control_topic = MQTT_CONFIG['control_topic']
feed_topic = MQTT_CONFIG['feed_topic']
device_topic = MQTT_CONFIG['device_topic']
username = MQTT_CONFIG['username']
password = MQTT_CONFIG['password']
client_id = f'activity-tracker-{random.randint(0, 1000)}'

# 截图配置
capture_interval = SCREENSHOT_CONFIG['capture_interval']

# 设备信息 - 使用用户名和机器名的哈希值作为设备ID的一部分，确保同一台电脑设备ID一致
username = getpass.getuser()
hostname = platform.node()
# 使用用户名和主机名生成固定的设备ID后缀
device_suffix = hashlib.md5(f"{username}_{hostname}".encode()).hexdigest()[:8]
device_id = f"{username}_{device_suffix}"
device_name = f"活动追踪器 {username}"

# 全局变量
camera = None
is_tracking = True
tracking_lock = threading.Lock()

# 截图设置
last_capture_time = 0

# MQTT 客户端
def on_connect(client, userdata, flags, rc):
    if rc == 0:
        print("✅ 活动追踪器已连接到 MQTT 服务器")
        # 发送上线通知
        send_device_status('online')
    else:
        print(f"❌ 连接失败，返回码 {rc}")

def send_device_status(status):
    """发送设备状态消息"""
    payload = {
        'device_id': device_id,
        'name': device_name,
        'status': status,
        'platform': platform.system()
    }
    client.publish(device_topic, json.dumps(payload))

def capture_screen():
    """截取屏幕截图"""
    try:
        with mss.mss() as sct:
            # 检查是否有可用的显示器
            if not sct.monitors:
                print("❌ 未检测到可用的显示器")
                return None
                
            # 获取主显示器（通常是monitors[1]，如果没有则使用monitors[0]）
            monitor = sct.monitors[1] if len(sct.monitors) > 1 else sct.monitors[0]
            
            # 获取屏幕截图
            screenshot = sct.grab(monitor)
            
            # 转换为numpy数组
            img = np.array(screenshot)
            
            # 检查图像是否有效
            if img is None or img.size == 0:
                print("❌ 屏幕截图为空")
                return None
                
            # 转换为BGR格式 (OpenCV使用BGR)
            # 首先确保图像是4通道的BGRA格式
            if len(img.shape) == 3 and img.shape[2] == 4:
                img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
            elif len(img.shape) == 3 and img.shape[2] == 3:
                # 如果已经是3通道的BGR格式，则不需要转换
                pass
            else:
                # 对于其他情况，尝试转换为BGR
                img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGR) if len(img.shape) == 2 else img
                
            return img
    except IndexError as e:
        print(f"❌ 显示器索引错误: {e}")
        return None
    except Exception as e:
        print(f"❌ 屏幕截图失败: {e}")
        return None

def on_message(client, userdata, msg):
    """处理MQTT消息"""
    try:
        payload = json.loads(msg.payload.decode())
        device_id_msg = payload.get('device_id')
        
        # 只处理发给本设备的消息
        if device_id_msg == device_id:
            command = payload.get('command')
            if command == 'start':
                global is_tracking
                is_tracking = True
                print("▶️ 收到开始指令")
            elif command == 'stop':
                is_tracking = False
                print("⏹️ 收到停止指令")
    except Exception as e:
        print(f"❌ 处理MQTT消息时出错: {e}")

def capture_camera():
    """捕获摄像头画面"""
    global camera
    try:
        if camera is None:
            camera = cv2.VideoCapture(0)
            camera.set(cv2.CAP_PROP_BUFFERSIZE, 1)
            camera.set(cv2.CAP_PROP_FPS, 10)
            camera.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
            camera.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
        
        ret, frame = camera.read()
        if ret:
            return frame
        else:
            print("❌ 无法读取摄像头画面")
            return None
    except Exception as e:
        print(f"❌ 摄像头捕获失败: {e}")
        return None

def send_image_to_mqtt(image, image_type):
    """将图像发送到MQTT"""
    try:
        # 编码图像
        _, buffer = cv2.imencode('.jpg', image, [cv2.IMWRITE_JPEG_QUALITY, 60])
        
        # 转为 base64 字符串
        jpg_as_text = base64.b64encode(buffer).decode('utf-8')
        
        # 创建包含图像类型和时间戳的消息
        payload = {
            'image': jpg_as_text,
            'type': image_type,  # 'screen' 或 'camera'
            'timestamp': datetime.now(pytz.timezone('Asia/Shanghai')).isoformat(),
            'device_id': device_id
        }
        
        # 发布到 MQTT
        result = client.publish(f"{feed_topic}/{device_id}", json.dumps(payload), qos=0)
        if result.rc != 0:
            print(f"❌ 发布失败: {result.rc}")
        else:
            print(f"📤 已发送{image_type}截图到服务器")
    except Exception as e:
        print(f"❌ 发送图像到MQTT失败: {e}")

def track_activity():
    """主跟踪函数"""
    global last_capture_time
    
    while is_tracking:
        try:
            current_time = time.time()
            if current_time - last_capture_time < capture_interval:
                time.sleep(0.01)
                continue
            
            last_capture_time = current_time
            
            # 捕获屏幕截图
            screen_img = capture_screen()
            if screen_img is not None:
                send_image_to_mqtt(screen_img, 'screen')
            
            # 捕获摄像头截图
            camera_img = capture_camera()
            if camera_img is not None:
                send_image_to_mqtt(camera_img, 'camera')
                
        except Exception as e:
            print(f"❌ 活动跟踪过程中出错: {e}")
            time.sleep(0.1)
    
    print("⏹️ 活动跟踪已停止")

def cleanup():
    """清理资源"""
    global camera, is_tracking
    is_tracking = False
    
    if camera and camera.isOpened():
        camera.release()
    
    # 发送离线通知
    send_device_status('offline')
    client.loop_stop()
    client.disconnect()
    print("👋 程序已关闭")

# 设置MQTT客户端
client = mqtt.Client(client_id)
client.username_pw_set(username, password)
client.on_connect = on_connect
client.on_message = on_message  # 添加消息回调

try:
    client.connect(broker, port, 60)
except Exception as e:
    print(f"❌ MQTT 连接失败: {e}")
    exit()

client.loop_start()

# 订阅控制主题
control_device_topic = f"{control_topic}/{device_id}"
client.subscribe(control_device_topic)
print(f"👂 已订阅控制主题: {control_device_topic}")

# 等待连接建立
time.sleep(1)

if __name__ == "__main__":
    try:
        print(f"📡 活动追踪器 {device_id} 已启动")
        print("📸 开始定时截图和AI分析...")
        
        # 启动活动跟踪线程
        tracking_thread = threading.Thread(target=track_activity, daemon=True)
        tracking_thread.start()
        
        # 保持程序运行
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        print("\n👋 正在关闭程序...")
        cleanup()